"""5.3【提示类】Prompt 标准化任务"""
import asyncio
import json
from contextlib import AsyncExitStack

from mcp import StdioServerParameters, stdio_client, ClientSession, ListResourcesResult, ListPromptsResult
from mcp.client.sse import sse_client
from openai import OpenAI

from deepseek_config import deepseek_v3_config


class MCPClient:



    def __init__(self):
        self.async_exit_stack = AsyncExitStack()
        self.session = None
        self.deepseek = OpenAI(
            api_key=deepseek_v3_config["api_key"],
            base_url=deepseek_v3_config["url"]
        )
        self.prompts = None


    async def connect_to_server(self,url:str):
        # 一 、创建stdio_client
        client = sse_client(url=url)
        transport = await self.async_exit_stack.enter_async_context(client)
        read_stream, write_stream = transport
        # 二、创建会话client
        client_session = ClientSession(read_stream, write_stream)
        self.session = await self.async_exit_stack.enter_async_context(client_session)
        # 三、初始化会话
        await self.session.initialize()

    async def execute(self,question:str):
        # 一、获取server服务端中的提示类list_prompt列表 (需要改动）
        list_prompt:ListPromptsResult = await self.session.list_prompts()
        prompts = list_prompt.prompts
        print("打印出获取的工具列表：",prompts)
        self.prompts = {}
        # 二、创建function calling 格式(大模型使用）、
        prompts_tools = []
        for resource in prompts:
            name = resource.name
            description = resource.description
            prompts_tools.append(
                {
                    "type": "function",
                    "function": {
                        "name": name,
                        "description": description,
                        "parameters": None
                    }
                }
            )
            #初始化resources属性
            self.prompts[name] = {
                "name":name,
                "description":description,
                "mime_type":None
            }

        # 三、 创建messages,deepseek大模型的格式
        messages = [
            {
                "role":"user",
                "content":question
            }
        ]
        # 四、调用deepseek大模型
        deepseek_response = self.deepseek.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            tools=prompts_tools
        )
        # 打印出大模型的决策结果
        print("==== deepseek 响应持结果：",deepseek_response)
        choice_result = deepseek_response.choices[0]
        #第二次调用大模型的前置参数
        messages.append(choice_result.message.model_dump())
        tool_call = choice_result.message.tool_calls[0]

        print(" tool_call:",tool_call)
        print("大模型决策的最终结果,工具名称：",tool_call.function.name,",参数：",tool_call.function.arguments)
        function_name = tool_call.function.name
        arguments = json.loads(tool_call.function.arguments)
        # 五、调用工具链 （需要改动）
        prompts_result = await self.session.get_prompt(
            name=function_name,
            arguments=arguments
        )
        print("==== 工具调用结果：",prompts_result)
        #最终的结果
        result = prompts_result.messages[0].content.text
        print("==== 最终的结果：",result)

        # 六、使用大模型生成最终的结果，并且使用语言模型生成最终的结果
        messages.append({
            "role": "tool",
            "content": result,
            "tool_call_id": tool_call.id
        })
        # 再次调用大模型
        deepseek_response = self.deepseek.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            tools=prompts_tools,
        )
        # 获取最终的结果
        result = deepseek_response.choices[0].message.content
        print("==== 最终的结果：", result)


    #关闭资源
    async def aclose(self):
        await self.async_exit_stack.aclose()

async def main():
    client = MCPClient()
    try:
        await client.connect_to_server("http://localhost:8000/sse")
        while True:
            question = input("请输入问题：")
            if question == "exit":
                break
            await client.execute(question)

    except Exception as e:
        print(f"连接失败: {e}")
        return
    finally:
        await client.aclose()

if __name__ == "__main__":
    asyncio.run(main())